# -*- coding: utf-8 -*-

import os
import re
import sys
import ctypes
import time

try:
    import gdb
except ImportError:
    print("[-] gdb_alloc_trace cannot run standalone")
    sys.exit(0)

class trace_writter:
    def __init__(self, pid):
        # tracer: nop
        #
        #           <...>-7844  (-----) [000] ...1   886.340626: tracing_mark_write: B|3073|queueBuffer
        #           <...>-4061  (-----) [000] ...1   886.340908: tracing_mark_write: B|3049|queueBuffer
        #           <...>-4061  (-----) [000] ...1   886.340932: tracing_mark_write: B|3049|SurfaceView: 2

        self.pid = pid
        #self.exe = gdb.current_progspace().filename
        self.f = open("/data/gdb_malloc.trace", "w")
        self.f.write("# tracer: malloc\n#\n")
#        self.f.write("# entries-in-buffer/entries-written: 1701/1701   #P:2")
        self.f.write("#\n")
        self.f.write("#                                      _-----=> irqs-off\n")
        self.f.write("#                                     / _----=> need-resched\n")
        self.f.write("#                                    | / _---=> hardirq/softirq\n")
        self.f.write("#                                    || / _--=> preempt-depth\n")
        self.f.write("#                                    ||| /     delay\n")
        self.f.write("#           TASK-PID    TGID   CPU#  ||||    TIMESTAMP  FUNCTION\n")
        self.f.write("#              | |        |      |   ||||       |         |\n")
        self.start_record = False

    def _trace_write(self, tid, time, str):
        self.f.write("     <...>-{:}  (-----) [000] ...1   {:0.6f}: tracing_mark_write: ".format(self.pid, time))
        self.f.write(str)
        self.f.write("\n")
        self.f.flush()

    def add_begin(self, tid, time, ptr, size):
        self._trace_write(tid, time, "B|{:}|0x{:x} {:}\n".format(self.pid, ptr, size))
        self.add_int(time, tid, "{:x}.format(ptr)", size)
    def add_end(self, tid, time):
        self._trace_write(tid, time, "E\n")

    def add_begin_async(self, tid, time, ptr, size):
        self._trace_write(tid, time, "S|{:}|alloc|0x{:x}".format(tid, ptr))
        self.add_int(time, tid, "{:x}".format(ptr), size)
    def add_end_async(self, tid, time, ptr):
        self._trace_write(tid, time, "F|{:}|free|0x{:x}".format(tid, ptr))

    def add_int(self, time, tid, n,v):
        self._trace_write(tid, time, "C|{:}|{:}|{:}".format(self.pid, n, v))

    def record_total(self, time, tid, size):
        if (self.start_record):
            self._trace_write(0, time, "B|total|{:}\n".format(self.pid, size))
        self._trace_write(0, time, "E\n".format(self.pid))
        self.add_int(time, tid, "total", size)

class proc_memory_info():
    def __init__(self, pid):
        self.statm_path = "/proc/{:}/statm".format(pid)
        self.statm_info = { 'total' : 0, 'anon' : 0, 'file' : 0, 'share' : 0 }

    def get_rss(self):
        with open(self.statm_path, 'r') as f:
            data = f.read()
            f.close()
            tmp = data.split()

            self.statm_info['total'] = int(tmp[1]) * 4
            self.statm_info['file'] = int(tmp[2]) * 4
            self.statm_info['share'] = int(tmp[3]) * 4
            self.statm_info['anon'] = self.statm_info['total'] - self.statm_info['file'] - self.statm_info['share']
        return self.statm_info

class proc_memory_smaps():
    def __init__(self, pid):
        self.smaps_file = "/proc/{:}/smaps".format(pid)
        self.smaps_info = { 'total' : 0, 'anon' : 0, 'file' : 0, 'share' : 0 }

    def get_rss(self):
        mem_type = ""
        self.smaps_info = { 'total' : 0, 'anon' : 0, 'file' : 0, 'share' : 0 }
        with open(self.smaps_file, 'r') as f:
            for line in f:
                vre = re.compile('^[0-9a-f]')
                res = vre.match(line)
                if res:
                    tmp = line.split()
                    if (len(tmp) <= 5):
                        mem_type = 'anon'
                    else:
                        if tmp[5].startswith("/"):
                            mem_type = 'file'
                        else:
                            mem_type = 'anon'
                else:
                    if line.startswith('Rss:') or line.startswith('Swap:'):
                        size = int(line.split()[1])
                        self.smaps_info[mem_type] += size
                        self.smaps_info['total'] += size
                    elif line.startswith('Shared_'):
                        self.smaps_info['share'] += int(line.split()[1])
        return self.smaps_info

class collect_memory_info(proc_memory_smaps):
    def __init__(self, pid, f, starttime):
        proc_memory_smaps.__init__(self, pid);
        self.f = f
        self.starttime = starttime
    def collect_rss(self):
        statm_info = self.get_rss()
        t = time.time() - self.starttime
        str = "{:0.6f}: {:} {:} {:} {:}".format(
            t,
            statm_info['total'],
            statm_info['anon'],
            statm_info['file'],
            statm_info['share'])
        self.f.write(str)
        self.f.write("\n")

        global tracer
        tracer.record_total(t, 0, statm_info['total'])
        print(str)

def get_current_thread_id():
    current_thread_info = gdb.execute("thread", to_string=True)
    idx = current_thread_info.find("(LWP ")
    if idx == -1:
        return -1
    idx2 = current_thread_info.find(")]")
    if idx2 == -1:
        return -1

    return int(current_thread_info[idx + 5 : idx2])

class tickle_check_meminfo():
    def __init__(self, max_tick, callback):
        self.max_tick = max_tick
        self.callback = callback
        self.tick = 0
        self.callback()
    def inc_and_trigger(self):
        self.tick += 1
        if self.tick >= self.max_tick:
            self.callback()
            self.tick = 0

def log_allocation(f, time, ptr, size, trace):
    global tracer
    global tickler
    tid = get_current_thread_id()
    str = "{:0.6f}: M {:} 0x{:x} {:}".format(time, tid, ptr, size)

    f.write(str)
    f.write("\n")
    f.write(trace)
    f.flush()

    tickler.inc_and_trigger()

    if tracer:
        tracer.add_begin_async(tid, time, ptr, size)
    print(str)

def log_dallocation(f, time, ptr):
    global tracer
    global tickler
    tid = get_current_thread_id()
    str = "{:0.6f}: F {:} 0x{:x}".format(time, tid, ptr)

    f.write(str)
    f.write("\n")
    f.flush()

    tickler.inc_and_trigger()

    if tracer:
        tracer.add_end_async(tid, time, ptr)
    print(str)

def get_reg_val(name):
    """Return a register's value."""""
    value = gdb.parse_and_eval(name)
    if value.type.code == gdb.TYPE_CODE_INT:
        mask = (1 << 64) - 1
        return int(value.cast(gdb.Value(mask).type)) & mask
    else:
        return int(value)

    #return int(gdb.execute(cmd, to_string=True).split()[1].strip(), 16)

class MallocFinishBreakPoint(gdb.FinishBreakpoint):
    def __init__(self, f, reg, size, starttime):
        gdb.FinishBreakpoint.__init__(self, gdb.newest_frame(), internal=True)
        self.silent = True
        self.reg = reg
        self.f = f
        self.size = size
        self.starttime = starttime
        self.get_ptr_cmd = "$" + self.reg + "0"

    def stop(self):
        ptr = get_reg_val(self.get_ptr_cmd)
        if (ptr == 0):
            return False;

        backtrace = gdb.execute("bt 8", to_string=True)
        idx = backtrace.find("Backtrace stopped")
        if idx > 0:
            backtrace = backtrace[0:idx]
        log_allocation(self.f, time.time() - self.starttime, ptr, self.size, backtrace)
        return False

    def out_of_scope(self):
        print("MallocFinishBreakPoint out_of_scope")


class MallocBreakPoint(gdb.Breakpoint):
    def __init__(self, spec, f, time, reg):
        gdb.Breakpoint.__init__(self, spec)
        self.silent = True
        self.f = f
        self.starttime = starttime
        self.reg = reg
        self.get_size_cmd = "$" + self.reg + "4"
        # Dump of assembler code for function malloc:
        #    0xac2f2240 <+0>:push {r4, lr}
        #    0xac2f2242 <+2>:mov r4, r0
        #    0xac2f2244 <+4>:ldr r0, [pc, #52]; (0xac2f227c <malloc+60>)
        #    0xac2f2246 <+6>:add r0, pc
        #    0xac2f2248 <+8>:ldr r0, [r0, #36]; 0x24
        # => 0xac2f224a <+10>:dmb ish
        #    0xac2f224e <+14>:cbnz r0, 0xac2f2270 <malloc+48>

    def stop(self):
        size = get_reg_val(self.get_size_cmd)

        MallocFinishBreakPoint(self.f, self.reg, size, self.starttime)
        return False

class ReallocFinishBreakPoint(gdb.FinishBreakpoint):
    def __init__(self, f, reg, size, orgptr, starttime):
        gdb.FinishBreakpoint.__init__(self, gdb.newest_frame(), internal=True)
        self.silent = True
        self.reg = reg
        self.orgptr = orgptr
        self.f = f
        self.size = size
        self.starttime = starttime
        self.get_ptr_cmd = "$" + self.reg + "0"

    def stop(self):
        ptr = get_reg_val(self.get_ptr_cmd)
        if (ptr == 0):
            return False;

        backtrace = gdb.execute("bt 8", to_string=True)
        log_dallocation(self.f, time.time() - self.starttime, self.orgptr)
        log_allocation(self.f, time.time() - self.starttime, ptr, self.size, backtrace)
        return False

    def out_of_scope(self):
        print("ReallocFinishBreakPoint out_of_scope")

class ReallocBreakPoint(gdb.Breakpoint):
    def __init__(self, spec, f, time, reg):
        gdb.Breakpoint.__init__(self, spec)
        self.silent = True
        self.f = f
        self.starttime = starttime
        self.reg = reg
        self.get_orgptr_cmd = "$" + self.reg + "4"
        self.get_size_cmd = "$" + self.reg + "5"
        #    0xac2f2368 <+0>:push {r4, r5, r7, lr}
        #    0xac2f236a <+2>:sub sp, #8
        #    0xac2f236c <+4>:mov r4, r0
        #    0xac2f236e <+6>:ldr r0, [pc, #68]; (0xac2f23b4 <realloc+76>)
        #    0xac2f2370 <+8>:mov r5, r1
        #    0xac2f2372 <+10>:add r0, pc
        #    0xac2f2374 <+12>:ldr r0, [r0, #36]; 0x24
        #    0xac2f2376 <+14>:dmb ish
        #    0xac2f237a <+18>:cbnz r0, 0xac2f238c <realloc+36>

    def stop(self):
        size = get_reg_val(self.get_size_cmd)
        orgptr = get_reg_val(self.get_orgptr_cmd)

        ReallocFinishBreakPoint(self.f, self.reg, size, orgptr, self.starttime)
        return False

class FreeBreakPoint(gdb.Breakpoint):
    def __init__(self, spec, f, time, reg):
        gdb.Breakpoint.__init__(self, spec)
        self.silent = True
        self.f = f
        self.starttime = starttime
        self.reg = reg
        self.get_ptr_cmd ="$" + self.reg + "0"

    def stop(self):
        ptr = get_reg_val(self.get_ptr_cmd)
        if (ptr != 0):
            log_dallocation(self.f, time.time() - self.starttime, ptr)

        return False


tracer = None
tickler = None

if __name__  == "__main__":
    gdb_init_cmd = [
        "set confirm off",
        "set verbose off",
        "set pagination off",
        "set print thread-events off",
        "set step-mode on",
        "set history save on",
        "set history filename /data/gdb_history",
        "set output-radix 0x10",
        "set print pretty on",
        "handle SIGALRM nostop print nopass",
        "handle SIG33 nostop noprint nopass",
        "handle SIG34 nostop noprint nopass",
        "handle SIG35 nostop noprint nopass",
        "handle SIGQUIT nostop noprint nopass",
    ]

    for cmd in gdb_init_cmd:
        gdb.execute(cmd)

    arch = gdb.selected_frame().architecture()
    if (arch.name() == "aarch64"):
        regprefix = "x"
    elif (arch.name().startswith("arm")):
        regprefix = "r"
    else:
        raise RuntimeError("unsupported arch")

    pid = gdb.selected_inferior().pid

    f = open("/data/gdb_malloc.txt", "w")

    tracer = trace_writter(pid)

    starttime = time.time()

    collector = collect_memory_info(pid, f, starttime)
    tickler = tickle_check_meminfo(50, collector.collect_rss)

    MallocBreakPoint("malloc", f, starttime, regprefix)
    ReallocBreakPoint("realloc", f, starttime, regprefix)
    FreeBreakPoint("free", f, starttime, regprefix)

    gdb.execute("c")
